import matplotlib.pyplot as plt
import os

# 创建目录（如果不存在的话）
output_dir = 'plot'
os.makedirs(output_dir, exist_ok=True)
output_path = os.path.join(output_dir, 'training_num.png')

iterations = [40, 80, 120, 160, 200, 400, 600, 800, 1000]
iterations = [i*1.622 for i in iterations]
success_trajectory_number = [0.25, 0.27, 0.26, 0.28, 0.33, 0.32, 0.33, 0.34, 0.35]   # 示例数据，可以替换为你的数据
# step_level_preference_data_number = [1111, 1454, 1622] # 示例数据，可以替换为你的数据

# 创建图表和第一个y轴
fig, ax1 = plt.subplots(figsize=(10, 6))

# 绘制第一个y轴上的折线图，并为每个点添加不同的标记
ax1.plot(iterations, success_trajectory_number, 'o-', markersize=18, color='#EA8379', label='Success Trajectory', linewidth=3.5)
# ax1.plot(iterations, step_level_preference_data_number, 's-', markersize=15, color='#7DAEE0', label='Step-Level Preference Data')

# 设置坐标轴标签和颜色
ax1.set_xlabel('Number of Preference Data for Training', fontsize=26)
ax1.set_ylabel('Average Reward on Test Set', fontsize=26, color='black')  # 将纵轴颜色改为黑色
ax1.tick_params(axis='x', labelcolor='black', labelsize=22)  # 更改y轴标签颜色为黑色
ax1.tick_params(axis='y', labelcolor='black', labelsize=22)  # 更改y轴标签颜色为黑色
# ax1.set_ylim(500, 1700)  # 手动设置第一个y轴的范围
# ax1.set_xticks(range(min(iterations), max(iterations) + 1, 20))


# 添加图例
# ax1.legend(loc='upper left', fontsize=26)

# 美化作图
plt.grid(True, which='both', linestyle='--', linewidth=0.5)
plt.tight_layout()

# 保存图表
plt.savefig(output_path, dpi=300)
plt.show()
